{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tb8pA0kWrURv"
   },
   "source": [
    "# Migrating Monolithic Prompts to Command-R with RAG\n",
    "\n",
    "Command-R is a powerful LLM optimized for long context tasks such as retrieval augmented generation (RAG). Migrating a monolithic task such as question-answering or query-focused summarization to RAG can improve the quality of responses due to reduced hallucination and improved conciseness through grounding.\n",
    "\n",
    "Previously, migrating an existing use case to RAG involved a lot of manual work around indexing documents, implementing at least a basic search strategy, extensive post-processing to introduce proper grounding through citations, and of course fine-tuning an LLM to work well in the RAG paradigm.\n",
    "\n",
    "This cookbook demonstrates automatic migration of monolithic prompts through two diverse use cases where an original prompt is broken down into two parts: (1) context; and (2) instructions. The former can be done automatically or through simple chunking, while the latter is done automatically by Command-R through single shot prompt optimization.\n",
    "\n",
    "The two use cases demonstrated here are:\n",
    "\n",
    "1. Autobiography Assistant; and\n",
    "2. Legal Question Answering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3W6cgQxEq09t"
   },
   "outputs": [],
   "source": [
    "# Start by installing the Cohere API\n",
    "#!pip install cohere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Hwl6hOBy34Z"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import re\n",
    "\n",
    "import cohere\n",
    "import getpass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YyD4Okmsy-AU",
    "outputId": "1006c70f-9e87-48b5-d43f-766905937c3d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cohere API key:··········\n"
     ]
    }
   ],
   "source": [
    "CO_API_KEY = getpass.getpass('cohere API key:')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dqFUGLVPzETw"
   },
   "outputs": [],
   "source": [
    "co = cohere.Client(CO_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jY-SXbzS-ZFt"
   },
   "source": [
    "## Autobiography Assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JfS7y2zD-ehv"
   },
   "source": [
    "This application scenario is a common LLM-as-assistant use case. Given some context, help the user to complete a task. In this case, the task is to write a concise autobiographical summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bHTwH-ac-xKE"
   },
   "outputs": [],
   "source": [
    "original_prompt = '''## information\n",
    "Current Job Title: Senior Software Engineer\n",
    "Current Company Name: GlobalSolTech\n",
    "Work Experience: Over 15 years of experience in software engineering, specializing in AI and machine learning. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.\n",
    "Current Department Name: AI Research and Development\n",
    "Education: B.Sc. in Physics from Trent University (2004), Ph.D. in Statistics from HEC in Paris (2010)\n",
    "Hobbies: I love hiking in the mountains, free diving, and collecting and restoring vintage world war one mechanical watches.\n",
    "Family: Married with 4 children and 3 grandchildren.\n",
    "\n",
    "## instructions\n",
    "Your task is to assist a user in writing a short biography for social media.\n",
    "The length of the text should be no more than 100 words.\n",
    "Write the summary in first person.'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z4513hfd-yKT"
   },
   "outputs": [],
   "source": [
    "response = co.chat(\n",
    "    message=original_prompt,\n",
    "    model='command-r',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BOUodUWc-yN7",
    "outputId": "ce974241-289a-4da2-ecee-9997bff1b6db"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm a Senior Software Engineer at GlobalSolTech, with over 15 years of experience in AI and machine learning. My expertise lies in developing innovative algorithms for natural language processing, computer vision, and recommendation systems. I hold a B.Sc. in Physics and a Ph.D. in Statistics and enjoy hiking, free diving, and collecting vintage watches in my spare time. I'm passionate about using my skills to contribute to cutting-edge AI research and development. At GlobalSolTech, I'm proud to be part of a dynamic team driving technological advancement.\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vZv-cVPK_gQt"
   },
   "source": [
    "Using Command-R, we can automatically upgrade the original prompt to a RAG-style prompt to get more faithful adherence to the instructions, a clearer and more concise prompt, and in-line citations for free. Consider the following meta-prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X-ArivFl_G2D"
   },
   "outputs": [],
   "source": [
    "meta_prompt = f'''Below is a task for an LLM delimited with ## Original Task. Your task is to split that task into two parts: (1) the context; and (2) the instructions.\n",
    "The context should be split into several separate parts and returned as a JSON object where each part has a name describing its contents and the value is the contents itself.\n",
    "Make sure to include all of the context contained in the original task description and do not change its meaning.\n",
    "The instructions should be re-written so that they are very clear and concise. Do not change the meaning of the instructions or task, just make sure they are very direct and clear.\n",
    "Return everything in a JSON object with the following structure:\n",
    "\n",
    "{{\n",
    "  \"context\": [{{\"<description of part 1>\": \"<content of part 1>\"}}, ...],\n",
    "  \"instructions\": \"<the re-written instructions>\"\n",
    "}}\n",
    "\n",
    "## Original Task\n",
    "{original_prompt}\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iaAHkuOTGTp9",
    "outputId": "98f2b664-e2aa-4718-9510-a569723dd1f5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is a task for an LLM delimited with ## Original Task. Your task is to split that task into two parts: (1) the context; and (2) the instructions.\n",
      "The context should be split into several separate parts and returned as a JSON object where each part has a name describing its contents and the value is the contents itself.\n",
      "Make sure to include all of the context contained in the original task description and do not change its meaning.\n",
      "The instructions should be re-written so that they are very clear and concise. Do not change the meaning of the instructions or task, just make sure they are very direct and clear.\n",
      "Return everything in a JSON object with the following structure:\n",
      "\n",
      "{\n",
      "  \"context\": [{\"<description of part 1>\": \"<content of part 1>\"}, ...],\n",
      "  \"instructions\": \"<the re-written instructions>\"\n",
      "}\n",
      "\n",
      "## Original Task\n",
      "## information\n",
      "Current Job Title: Senior Software Engineer\n",
      "Current Company Name: GlobalSolTech\n",
      "Work Experience: Over 15 years of experience in software engineering, specializing in AI and machine learning. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.\n",
      "Current Department Name: AI Research and Development\n",
      "Education: B.Sc. in Physics from Trent University (2004), Ph.D. in Statistics from HEC in Paris (2010)\n",
      "Hobbies: I love hiking in the mountains, free diving, and collecting and restoring vintage world war one mechanical watches.\n",
      "Family: Married with 4 children and 3 grandchildren.\n",
      "\n",
      "## instructions\n",
      "Your task is to assist a user in writing a short biography for social media.\n",
      "The length of the text should be no more than 100 words.\n",
      "Write the summary in first person.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(meta_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4Ze_pzuD_xAz"
   },
   "source": [
    "Command-R returns with the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qxdiM_8I_HQF"
   },
   "outputs": [],
   "source": [
    "upgraded_prompt = co.chat(\n",
    "    message=meta_prompt,\n",
    "    model='command-r',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GEB3PEpy_HSs",
    "outputId": "e58cb699-a65f-4f80-d6bd-9c01b15d6127"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here is the task delved into a JSON object as requested:\n",
      "```json\n",
      "{\n",
      "  \"context\": [\n",
      "    {\n",
      "      \"Work Experience\": \"Over 15 years of AI and machine learning engineering experience. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.\"\n",
      "    },\n",
      "    {\n",
      "      \"Education\": \"B.Sc. in Physics (Trent University, 2004) and Ph.D. in Statistics (HEC Paris, 2010).\"\n",
      "    },\n",
      "    {\n",
      "      \"Personal Life\": \"I’m a married senior software engineer with 4 children and 3 grandchildren. I enjoy hiking, free diving, and vintage watch restoration.\"\n",
      "    },\n",
      "    {\n",
      "      \"Current Position\": \"I work at GlobalSolTech in the AI Research and Development department as a senior software engineer.\"\n",
      "    }\n",
      "  ],\n",
      "  \"instructions\": \"Using the provided information, write a concise, first-person social media biography of no more than 100 words.\"\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "print(upgraded_prompt.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gHHSfczbAAXL"
   },
   "source": [
    "To extract the returned information, we will write two simple functions to post-process out the JSON and then parse it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ebYeBAoj_HXS"
   },
   "outputs": [],
   "source": [
    "def get_json(text: str) -> str:\n",
    "    matches = [m.group(1) for m in re.finditer(\"```([\\w\\W]*?)```\", text)]\n",
    "    if len(matches):\n",
    "        postproced = matches[0]\n",
    "        if postproced[:4] == 'json':\n",
    "            return postproced[4:]\n",
    "        return postproced\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wSHiCxBr__2X"
   },
   "outputs": [],
   "source": [
    "def get_prompt_and_docs(text: str) -> tuple:\n",
    "    json_obj = json.loads(get_json(text))\n",
    "    prompt = json_obj['instructions']\n",
    "    docs = []\n",
    "    for item in json_obj['context']:\n",
    "        for k,v in item.items():\n",
    "            docs.append({\"title\": k, \"snippet\": v})\n",
    "    return prompt, docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mXLJbR8oAI-m"
   },
   "outputs": [],
   "source": [
    "new_prompt, docs = get_prompt_and_docs(upgraded_prompt.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rQEZ4ZcEAJMG",
    "outputId": "7554b49d-a544-45bc-bc9e-d080618a07bb"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('Using the provided information, write a concise, first-person social media biography of no more than 100 words.',\n",
       " [{'title': 'Work Experience',\n",
       "   'snippet': 'Over 15 years of AI and machine learning engineering experience. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.'},\n",
       "  {'title': 'Education',\n",
       "   'snippet': 'B.Sc. in Physics (Trent University, 2004) and Ph.D. in Statistics (HEC Paris, 2010).'},\n",
       "  {'title': 'Personal Life',\n",
       "   'snippet': 'I’m a married senior software engineer with 4 children and 3 grandchildren. I enjoy hiking, free diving, and vintage watch restoration.'},\n",
       "  {'title': 'Current Position',\n",
       "   'snippet': 'I work at GlobalSolTech in the AI Research and Development department as a senior software engineer.'}])"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_prompt, docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lpkkPFXCA95M"
   },
   "source": [
    "As we can see above, the new prompt is much more concise and gets right to the point. The context has been split into 4 \"documents\" that Command-R can ground the information to. Now let's run the same task with the new prompt while leveraging the `documents=` parameter. Note that the `docs` variable is a list of dict objects with `title` describing the contents of a text and `snippet` containing the text itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "40n-ziCtAJN7"
   },
   "outputs": [],
   "source": [
    "response = co.chat(\n",
    "    message=new_prompt,\n",
    "    model='command-r',\n",
    "    documents=docs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VdGi9H54AJP5",
    "outputId": "e6f3b651-b3e4-412d-8a62-636b119d9dba"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm a senior software engineer with a Ph.D. in Statistics and over 15 years of AI and machine learning engineering experience. My current focus at GlobalSolTech's AI R&D department is developing algorithms for natural language processing, computer vision, and recommendation systems. In my free time, I enjoy hiking, freediving, and restoring vintage watches, and I'm a married father of four with three grandchildren.\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Js3K4sLXBYW3"
   },
   "source": [
    "The response is concise. More importantly, we can ensure that there is no hallucination because the text is automatically grounded in the input documents. Using the simple function below, we can add this grounding information to the text as citations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n3UffjSNAJRs"
   },
   "outputs": [],
   "source": [
    "def insert_citations(text: str, citations: list[dict], add_one: bool=False):\n",
    "    \"\"\"\n",
    "    A helper function to pretty print citations.\n",
    "    \"\"\"\n",
    "    offset = 0\n",
    "    # Process citations in the order they were provided\n",
    "    for citation in citations:\n",
    "        # Adjust start/end with offset\n",
    "        start, end = citation.start + offset, citation.end + offset\n",
    "        if add_one:\n",
    "            cited_docs = [str(int(doc[4:]) + 1) for doc in citation.document_ids]\n",
    "        else:\n",
    "            cited_docs = [doc[4:] for doc in citation.document_ids]\n",
    "        # Shorten citations if they're too long for convenience\n",
    "        if len(cited_docs) > 3:\n",
    "            placeholder = \"[\" + \", \".join(cited_docs[:3]) + \"...]\"\n",
    "        else:\n",
    "            placeholder = \"[\" + \", \".join(cited_docs) + \"]\"\n",
    "        # ^ doc[4:] removes the 'doc_' prefix, and leaves the quoted document\n",
    "        modification = f'{text[start:end]} {placeholder}'\n",
    "        # Replace the cited text with its bolded version + placeholder\n",
    "        text = text[:start] + modification + text[end:]\n",
    "        # Update the offset for subsequent replacements\n",
    "        offset += len(modification) - (end - start)\n",
    "\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "n-8ZPQYVBtEt",
    "outputId": "c24ea218-1862-4463-be13-0f8cbad8e565"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm a senior software engineer [3, 4] with a Ph.D. in Statistics [2] and over 15 years of AI and machine learning engineering experience. [1] My current focus at GlobalSolTech's AI R&D department [4] is developing algorithms for natural language processing, computer vision, and recommendation systems. [1] In my free time, I enjoy hiking, freediving, and restoring vintage watches [3], and I'm a married father of four with three grandchildren. [3]\n"
     ]
    }
   ],
   "source": [
    "print(insert_citations(response.text, response.citations, True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yFDXXrf0CRJF"
   },
   "source": [
    "Now let's move on to an arguably more difficult problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qCl6ifD8tpBW"
   },
   "source": [
    "## Legal Question Answering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cEdYRffp06d2"
   },
   "source": [
    "On March 21st, the DOJ announced that it is [suing apple](https://www.theverge.com/2024/3/21/24107659/apple-doj-lawsuit-antitrust-documents-suing) for anti-competitive practices. The [complaint](https://www.justice.gov/opa/media/1344546/dl) is 88 pages long and consists of about 230 paragraphs of text. To understand what the suit alleges, a common use case would be to ask for a summary. Because Command-R has a context window of 128K, even an 88-page legal complaint fits comfortably within the window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ObQvkG-Aq_Aw"
   },
   "outputs": [],
   "source": [
    "apple = open('data/apple_mod.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5iBgO-XZtsWI",
    "outputId": "4af2276f-fac3-4d73-9723-e759a27f9f5b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "29697"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = co.tokenize(text=apple, model='command-r')\n",
    "len(tokens.tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LOfIr_O92cIe"
   },
   "source": [
    "We can set up a prompt template that allows us to ask questions on the original text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "csTeS7zx2PCY"
   },
   "outputs": [],
   "source": [
    "prompt_template = '''\n",
    "# Legal Text\n",
    "{legal_text}\n",
    "\n",
    "# Question\n",
    "{question}\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Hk_41bi2la5"
   },
   "outputs": [],
   "source": [
    "question = '''Please summarize the attached legal complaint succinctly. Focus on answering the question: what does the complaint allege?'''\n",
    "rendered_prompt = prompt_template.format(legal_text=apple, question=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uGEe_X7H2mI9"
   },
   "outputs": [],
   "source": [
    "response = co.chat(\n",
    "    message=rendered_prompt,\n",
    "    model='command-r',\n",
    "    temperature=0.3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JLRlSA5_2mLb",
    "outputId": "63a6cc71-19d6-4382-a1c1-7aa14c889ef6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The complaint alleges that Apple has violated antitrust laws by engaging in a pattern of anticompetitive conduct to maintain its monopoly power over the U.S. markets for smartphones and performance smartphones. Apple is accused of using its control over app distribution and access to its operating system to impede competition and innovation. Specifically, the company is said to have restricted developers' ability to create certain apps and limited the functionality of others, making it harder for consumers to switch away from iPhones to rival smartphones. This conduct is alleged to have harmed consumers and developers by reducing choice, increasing prices, and stifling innovation. The plaintiffs seek injunctive relief and potential monetary awards to remedy these illegal practices.\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7a81dU5U3LIY"
   },
   "source": [
    "The summary seems clear enough. But I am interested in the specific allegations that the DOJ makes. For example, skimming the full complaint, it looks like the DOJ is alleging that Apple could encrypt text messages sent to Android phones if it wanted to do so. We can ammend the rendered prompt and ask:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0DjeUClB2mOW"
   },
   "outputs": [],
   "source": [
    "question = '''Does the DOJ allege that Apple could encrypt text messages sent to Android phones?'''\n",
    "rendered_prompt = prompt_template.format(legal_text=apple, question=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1nRvWhRj88Io"
   },
   "outputs": [],
   "source": [
    "response = co.chat(\n",
    "    message=rendered_prompt,\n",
    "    model='command-r',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Cvp8jQDt88Qi",
    "outputId": "95db3931-fa98-4188-9481-b7e574d036d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, the DOJ alleges that Apple could allow iPhone users to send encrypted messages to Android users while still using iMessage on their iPhones but chooses not to do so. According to the DOJ, this would instantly improve the privacy and security of iPhones and other smartphones.\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wTqGg4rw90-y"
   },
   "source": [
    "This is a very interesting allegation that at first glance suggests that the model could be hallucinating. Because RAG has been shown to help reduce hallucinations and grounds its responses in the input text, we should convert this prompt to the RAG style paradigm to gain confidence in its response.\n",
    "\n",
    "While previously we asked Command-R to chunk the text for us, the legal complaint is highly structured with numbered paragraphs so we can use the following function to break the complaint into input docs ready for RAG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QztSRuzytuS7"
   },
   "outputs": [],
   "source": [
    "def chunk_doc(input_doc: str) -> list:\n",
    "    chunks = []\n",
    "    current_para = 'Preamble'\n",
    "    current_chunk = ''\n",
    "    # pattern to find an integer number followed by a dot (finding the explicitly numbered paragraph numbers)\n",
    "    pattern = r'^\\d+\\.$'\n",
    "\n",
    "    for line in input_doc.splitlines():\n",
    "        if re.match(pattern, line):\n",
    "            chunks.append((current_para.replace('.', ''), current_chunk))\n",
    "            current_chunk = ''\n",
    "            current_para = line\n",
    "        else:\n",
    "            current_chunk += line + '\\n'\n",
    "\n",
    "    docs = []\n",
    "    for chunk in chunks:\n",
    "        docs.append({\"title\": chunk[0], \"snippet\": chunk[1]})\n",
    "\n",
    "    return docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qDtlhxtIDeaD"
   },
   "outputs": [],
   "source": [
    "chunks = chunk_doc(apple)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jTya5ti9DizD",
    "outputId": "b545d55a-146b-47b5-9e6b-8173e8325141"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'title': '18', 'snippet': '\\nProtecting competition and the innovation that competition inevitably ushers in\\nfor consumers, developers, publishers, content creators, and device manufacturers is why\\nPlaintiffs bring this lawsuit under Section 2 of the Sherman Act to challenge Apple’s\\nmaintenance of its monopoly over smartphone markets, which affect hundreds of millions of\\nAmericans every day. Plaintiffs bring this case to rid smartphone markets of Apple’s\\nmonopolization and exclusionary conduct and to ensure that the next generation of innovators\\ncan upend the technological world as we know it with new and transformative technologies.\\n\\n\\nII.\\n\\nDefendant Apple\\n\\n'}\n"
     ]
    }
   ],
   "source": [
    "# example: show the 18th chunk\n",
    "print(chunks[18])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-ouAZC-iD1-o"
   },
   "source": [
    "We can now try the same question but ask it directly to Command-R with the chunks as grounding information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IILyU94CDlQu"
   },
   "outputs": [],
   "source": [
    "response = co.chat(\n",
    "    message='''Does the DOJ allege that Apple could encrypt text messages sent to Android phones?''',\n",
    "    model='command-r',\n",
    "    documents=chunks,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yJclGbq2D-38",
    "outputId": "10f787a4-9c15-455b-eef2-16645a211c77"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, according to the DOJ, Apple could encrypt text messages sent from iPhones to Android phones. The DOJ claims that Apple degrades the security and privacy of its users by impeding cross-platform encryption and preventing developers from fixing the broken cross-platform messaging experience. Apple's conduct makes it harder to switch from iPhone to Android, as messages sent from iPhones to Android phones are unencrypted.\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3qIa3ALlET3S"
   },
   "source": [
    "The responses seem similar, but we should add citations and check the citation to get confidence in the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WLp2RzU4EO7B",
    "outputId": "38266cbf-1fc7-47aa-b1e5-e727e077c331"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, according to the DOJ, Apple could encrypt text messages sent from iPhones to Android phones. [144] The DOJ claims that Apple degrades the security and privacy [144] of its users by impeding cross-platform encryption [144] and preventing developers from fixing the broken cross-platform messaging experience. [93] Apple's conduct makes it harder to switch from iPhone to Android [144], as messages sent from iPhones to Android phones are unencrypted. [144]\n"
     ]
    }
   ],
   "source": [
    "print(insert_citations(response.text, response.citations))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bBPZqjjFE_4f"
   },
   "source": [
    "The most important passage seems to be paragraph 144. Paragraph 93 is also cited. Let's check what they contain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6AM_xdHQEdvW",
    "outputId": "1e17f5d6-ebc4-4c90-f41c-48d4b766fb82"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Apple is also willing to make the iPhone less secure and less private if that helps\n",
      "maintain its monopoly power. For example, text messages sent from iPhones to Android phones\n",
      "are unencrypted as a result of Apple’s conduct. If Apple wanted to, Apple could allow iPhone\n",
      "users to send encrypted messages to Android users while still using iMessage on their iPhone,\n",
      "which would instantly improve the privacy and security of iPhone and other smartphone users.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(chunks[144]['snippet'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jDBfehaJFIl5",
    "outputId": "e20355ae-19b2-4d66-a013-97a431805d88"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Recently, Apple blocked a third-party developer from fixing the broken cross-\n",
      "platform messaging experience in Apple Messages and providing end-to-end encryption for\n",
      "messages between Apple Messages and Android users. By rejecting solutions that would allow\n",
      "for cross-platform encryption, Apple continues to make iPhone users’ less secure than they could\n",
      "otherwise be.\n",
      "\n",
      "ii.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(chunks[93]['snippet'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8R8EeVf5FSR2"
   },
   "source": [
    "Paragraph 144 indeed contains the important allegation: **If Apple wanted to, Apple could allow iPhone users to send encrypted messages to Android users**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DCt5OR2dHOWS"
   },
   "source": [
    "# Conclusion\n",
    "\n",
    "In this cookbook we have shown how one can easily take an existing monolithic prompt and migrate it to the RAG paradigm to get less hallucination, grounded information, and in-line citations. We also demonstrated Command-R's ability to re-write an instruction prompt in a single shot to make it more concise and potentially lead to higher quality completions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0vTsJc-ZHleg"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
